package org.apache.lucene.distributed;

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.ObjectStreamClass;
import java.io.OutputStream;
import java.io.Serializable;
import java.net.InetSocketAddress;
import java.rmi.Remote;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.commons.io.IOUtils;
import org.apache.lucene.distributed.RMI.ProxyKey;
import org.apache.lucene.distributed.RMI.Server.InternalService;

// TODO: need to load classes from jar names of classes that are in Lucene
// and store as a text file
// TODO: if class cannot be loaded from
/**
 * Offers object streaming with dynamic class loading.
 * 
 * The ClassKey holds the class name and serialVersionUID. The combination makes
 * a class unique. This allows classes that have changed to be loading
 * dynamically on deserialization in RMIObjectInputStream. This is because when
 * a class changes, the serialVersionUID changes as well.  Each class
 * is given a SingleClassLoader so that it can different versions of the class
 * may later be loaded and used.
 * 
 * RMIObjectOutputStream records classes that are ok for ClassAccept.isSpecial
 * in the classLoaders map. In the Lucene use case these would be subclasses of
 * for example the Query class serialized on the client side.
 * 
 * RMIObjectInputStream dynamically loads classes that are ClassAccept.isSpecial
 * from the the given CLInterface. If the class is new, a SingleClassLoader is
 * created and placed into classLoaders.
 * 
 */
public class RMIClasses implements CLInterface {
  ConcurrentHashMap<ClassKey,ClassLoader> classLoaders = new ConcurrentHashMap<ClassKey,ClassLoader>();
  ClassAccept classAccept;

  public RMIClasses(ClassAccept classAccept) {
    this.classAccept = classAccept;
  }
  
  /**
   * If a class is special is it will be dynamically loaded and given 
   * a SingleClassLoader.
   */
  public static interface ClassAccept {
    public boolean isSpecial(String name);
  }

  public byte[] getResource(String resource) throws IOException {
    throw new UnsupportedOperationException("");
  }

  public static String classToResource(String name) {
    return name.replace('.', '/') + ".class";
  }

  public byte[] loadClassData(ClassKey classKey) throws IOException {
    ClassLoader classLoader = classLoaders.get(classKey);
    if (classLoader != null) {
      String resource = classToResource(classKey.name);
      InputStream input = classLoader.getResourceAsStream(resource);
      return IOUtils.toByteArray(input);
    } else
      return null;
  }
  
  public static Class getRemoteInterface(Class cl) {
    Class[] is = cl.getInterfaces();
    for (Class i : is) {
      if (Remote.class.isAssignableFrom(i)) {
        return i;
      }
    }
    return null;
  }

  /**
   * Maps a ClassKey to a ClassLoader. ClassLoader will only load the Class
   * given by the ClassKey. The class bytes are read from CLInterface.
   * 
   */
  public static class SingleClassLoader extends ClassLoader {
    CLInterface clInterface;
    ClassKey classKey;

    public SingleClassLoader(ClassKey classKey, ClassLoader parent, CLInterface clInterface) {
      super(parent);
      this.classKey = classKey;
      this.clInterface = clInterface;
    }

    public InputStream getResourceAsStream(String name) {
      try {
        return new ByteArrayInputStream(clInterface.getResource(name));
      } catch (IOException ioException) {
        throw new RuntimeException(ioException);
      }
    }

    protected Class<?> findClass(final String name) throws ClassNotFoundException {
      if (name.equals(classKey.name)) {
        byte[] classData = null;
        try {
          classData = clInterface.loadClassData(classKey);
        } catch (Exception exception) {
          throw new ClassNotFoundException("", exception);
        }
        if (classData == null) {
          throw new ClassNotFoundException(name);
          //return super.findClass(name);
        }
        return defineClass(name, classData, 0, classData.length);
      } else {
        return super.findClass(name);
      }
    }

    protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
      Class c = findLoadedClass(name);
      if (c == null) {
        if (name.equals(classKey.name)) {
          c = findClass(name);
          if (resolve) {
            resolveClass(c);
          }
          return c;
        } else {
          return super.loadClass(name, resolve);
        }
      }
      return c;
    }
  }

  public class RMIObjectOutputStream extends ObjectOutputStream {
    private RMI.Server rmiServer;
    private boolean writeClass;

    public RMIObjectOutputStream(OutputStream output, RMI.Server rmiServer, boolean writeClass) throws IOException {
      super(output);
      this.rmiServer = rmiServer;
      this.writeClass = writeClass;
      if (rmiServer != null) {
        enableReplaceObject(true);
      }
    }

    /**
     * Overridden to write out proxy key in place of service object if needed
     */
    protected Object replaceObject(Object obj) throws IOException {
      assert rmiServer != null;
      String name = rmiServer.getServiceName(obj);
      if (name != null) {
        RMI rmi = rmiServer.getRMI();
        Class remoteInterface = getRemoteInterface(obj.getClass());
        InternalService is = rmiServer.getInternalService(name);
        return new ProxyKey(name, remoteInterface, new InetSocketAddress(rmi.localhost, rmi.port), is.dgc);
      } else {
        return obj;
      }
    }

    /**
     * Marks class as one that may be loaded by the receiver if it's special
     */
    protected void writeClassDescriptor(ObjectStreamClass desc) throws IOException {
      String name = desc.getName();
      if (classAccept.isSpecial(name)) {
        Class cl = desc.forClass();
        ClassKey classKey = new ClassKey(desc.getName(), desc.getSerialVersionUID());
        if (!classLoaders.containsKey(classKey) && !cl.isPrimitive()) {
          ClassLoader classLoader = cl.getClassLoader();
          if (classLoader == null) {
            assert classLoader != null;
          }
          classLoaders.put(classKey, classLoader);
          if (writeClass) {
            String resource = classToResource(name);
            byte[] bytes = IOUtils.toByteArray(classLoader.getResourceAsStream(name));
            writeInt(bytes.length);
            write(bytes);
          }
        }
      }
      super.writeClassDescriptor(desc);
    }
  }

  public RMIObjectOutputStream createOutput(OutputStream output) throws IOException {
    return createOutput(output, null);
  }

  public RMIObjectOutputStream createOutput(OutputStream output, RMI.Server rmiServer) throws IOException {
    return new RMIObjectOutputStream(output, rmiServer, false);
  }
  
  public RMIObjectOutputStream createOutput(OutputStream output, RMI.Server rmiServer, boolean writeClasses) throws IOException {
    return new RMIObjectOutputStream(output, rmiServer, writeClasses);
  }

  /**
   * Used by the the RMIClient to insure that serialized proxies are handled
   * 
   * @param input
   * @param rmi
   * @return
   * @throws IOException
   */
  public RMIObjectInputStream createInput(InputStream input, RMI rmi) throws IOException {
    return new RMIObjectInputStream(input, null, rmi, false);
  }
  
  public RMIObjectInputStream createInput(InputStream input, RMI rmi, boolean readClasses) throws IOException {
    return new RMIObjectInputStream(input, null, rmi, readClasses);
  }

  public RMIObjectInputStream createInput(InputStream input, CLInterface clInterface, RMI rmi) throws IOException {
    return new RMIObjectInputStream(input, clInterface, rmi, false);
  }
  
  /**
   * An ObjectInputStream that dynamically loads classes from the
   * CLInterface if given.  Also transforms proxy objects into
   * ProxyKey objects.
   *
   */
  public class RMIObjectInputStream extends ObjectInputStream {
    CLInterface clInterface;
    RMI rmi;
    boolean readClass;
    
    /**
     * 
     * @param input
     * @param clInterface
     * @param rmi
     * @param readClass If true class bytes are written to the serialized stream
     * @throws IOException
     */
    public RMIObjectInputStream(InputStream input, CLInterface clInterface, RMI rmi, boolean readClass) throws IOException {
      super(input);
      this.clInterface = clInterface;
      this.rmi = rmi;
      this.readClass = readClass;
      if (rmi != null) {
        enableResolveObject(true);
      }
    }

    protected Object resolveObject(Object obj) throws IOException {
      if (obj instanceof ProxyKey) {
        // transform to local rmi proxy
        return rmi.getProxy((ProxyKey) obj);
      } else {
        return obj;
      }
    }

    protected Class<?> resolveClass(ObjectStreamClass desc) throws IOException, ClassNotFoundException {
      String name = desc.getName();
      if (!classAccept.isSpecial(name)) {
        try {
          return super.resolveClass(desc);
        } catch (ClassNotFoundException ex) {
          // could not find class locally, try loading it
          throw new ClassNotFoundException(name);
        }
      }
      CLInterface cli = clInterface;
      long serialUID = desc.getSerialVersionUID();
      if (readClass) {
        int l = readInt();
        final byte[] classBytes = new byte[l];
        readFully(classBytes);
        cli = new CLInterface() {
          public byte[] getResource(String resource) throws IOException {
            return classBytes;
          }
          public byte[] loadClassData(ClassKey classKey) throws IOException {
            return classBytes;
          }
        };
      }
      ClassLoader classLoader = getClassLoader(new ClassKey(name, serialUID), cli);
      if (classLoader == null)
        return super.resolveClass(desc);
      Class c = classLoader.loadClass(name);
      if (c == null) {
        throw new ClassNotFoundException(name);
      }
      return c;
    }
  }

  public static class ClassKey implements Serializable {
    public String name;
    public Long serialUID;

    public ClassKey(String name, Long serialUID) {
      this.name = name;
      this.serialUID = serialUID;
    }
    
    public String toString() {
      return name+" "+serialUID;
    }
    
    @Override
    public int hashCode() {
      final int prime = 31;
      int result = 1;
      result = prime * result + ((name == null) ? 0 : name.hashCode());
      result = prime * result + ((serialUID == null) ? 0 : serialUID.hashCode());
      return result;
    }

    @Override
    public boolean equals(Object obj) {
      if (this == obj)
        return true;
      if (obj == null)
        return false;
      if (getClass() != obj.getClass())
        return false;
      final ClassKey other = (ClassKey) obj;
      if (name == null) {
        if (other.name != null)
          return false;
      } else if (!name.equals(other.name))
        return false;
      if (serialUID == null) {
        if (other.serialUID != null)
          return false;
      } else if (!serialUID.equals(other.serialUID))
        return false;
      return true;
    }

    
  }

  public ClassLoader getClassLoader(ClassKey classKey, CLInterface clInterface) {
    ClassLoader classLoader = classLoaders.get(classKey);
    if (classLoader == null) {
      ClassLoader parent = Thread.currentThread().getContextClassLoader();
      classLoader = new SingleClassLoader(classKey, parent, clInterface);
      classLoaders.put(classKey, classLoader);
    }
    return classLoader;
  }
}
